import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from collections import deque, namedtuple
import random

# 定义经验回放的元组结构
Experience = namedtuple('Experience', ['state', 'action', 'reward', 'next_state', 'done'])


class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def add(self, state, action, reward, next_state, done):
        self.buffer.append(Experience(state, action, reward, next_state, done))

    def sample(self, batch_size):
        experiences = random.sample(self.buffer, batch_size)
        return experiences

    def __len__(self):
        return len(self.buffer)


class DDPG:
    def __init__(self, state_dim, action_dim, action_bound, learning_rate=3e-4, gamma=0.99):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.action_bound = action_bound
        self.learning_rate = learning_rate
        self.gamma = gamma

        # 初始化演员网络和评论家网络
        self.actor = Actor(state_dim, action_dim, action_bound)
        self.critic = Critic(state_dim, action_dim)
        self.target_actor = Actor(state_dim, action_dim, action_bound)
        self.target_critic = Critic(state_dim, action_dim)

        # 复制网络权重到目标网络
        self.target_actor.load_state_dict(self.actor.state_dict())
        self.target_critic.load_state_dict(self.critic.state_dict())

        # 定义优化器
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=learning_rate)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=learning_rate)

        # 经验回放缓冲区
        self.memory = ReplayBuffer(10000)

    def select_action(self, state):
        state = torch.FloatTensor(state.reshape(1, -1))
        action = self.actor(state)
        return action.detach().numpy().flatten()

    def update(self, batch_size):
        # 从经验回放缓冲区中采样
        experiences = self.memory.sample(batch_size)
        states = torch.FloatTensor(np.vstack([e.state for e in experiences if e is not None]))
        actions = torch.FloatTensor(np.vstack([e.action for e in experiences if e is not None]))
        rewards = torch.FloatTensor(np.vstack([e.reward for e in experiences if e is not None]))
        next_states = torch.FloatTensor(np.vstack([e.next_state for e in experiences if e is not None]))
        dones = torch.FloatTensor(np.vstack([e.done for e in experiences if e is not None]))

        # 计算目标 Q 值
        next_actions = self.target_actor(next_states)
        Q_targets_next = self.target_critic(next_states, next_actions)
        Q_targets = rewards + (self.gamma * Q_targets_next * (1 - dones))

        # 更新评论家网络
        Q_expected = self.critic(states, actions)
        critic_loss = nn.MSELoss()(Q_expected, Q_targets)
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        # 更新演员网络
        actions_pred = self.actor(states)
        actor_loss = -self.critic(states, actions_pred).mean()
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # 软更新目标网络
        self.soft_update(self.critic, self.target_critic)
        self.soft_update(self.actor, self.target_actor)

    def soft_update(self, local_model, target_model, tau=1e-3):
        for target_param, local_param in zip(target_model.parameters(), local_model.parameters()):
            target_param.data.copy_(tau * local_param.data + (1.0 - tau) * target_param.data)


class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, action_bound):
        super(Actor, self).__init__()
        self.fc1 = nn.Linear(state_dim, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, action_dim)
        self.action_bound = action_bound

    def forward(self, state):
        x = torch.relu(self.fc1(state))
        x = torch.relu(self.fc2(x))
        x = torch.tanh(self.fc3(x)) * self.action_bound
        return x


class Critic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(state_dim, 256)
        self.fc2 = nn.Linear(256 + action_dim, 128)
        self.fc3 = nn.Linear(128, 1)

    def forward(self, state, action):
        x = torch.relu(self.fc1(state))
        x = torch.cat((x, action), dim=1)
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class DQN:
    def __init__(self, state_dim, action_dim, learning_rate=3e-4, gamma=0.99):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.learning_rate = learning_rate
        self.gamma = gamma

        # 初始化 Q 网络和目标网络
        self.q_network = QNetwork(state_dim, action_dim)
        self.target_q_network = QNetwork(state_dim, action_dim)
        self.target_q_network.load_state_dict(self.q_network.state_dict())

        # 定义优化器
        self.optimizer = optim.Adam(self.q_network.parameters(), lr=learning_rate)

        # 经验回放缓冲区
        self.memory = ReplayBuffer(10000)

    def select_action(self, state, epsilon=0.1):
        state = torch.FloatTensor(state.reshape(1, -1))
        if np.random.uniform() < epsilon:
            action = np.random.choice(self.action_dim)
        else:
            with torch.no_grad():
                q_values = self.q_network(state)
                action = torch.argmax(q_values, dim=1).numpy()[0]
        return action

    def update(self, batch_size):
        # 从经验回放缓冲区中采样
        experiences = self.memory.sample(batch_size)
        states = torch.FloatTensor(np.vstack([e.state for e in experiences if e is not None]))
        actions = torch.LongTensor(np.vstack([e.action for e in experiences if e is not None]))
        rewards = torch.FloatTensor(np.vstack([e.reward for e in experiences if e is not None]))
        next_states = torch.FloatTensor(np.vstack([e.next_state for e in experiences if e is not None]))
        dones = torch.FloatTensor(np.vstack([e.done for e in experiences if e is not None]))

        # 计算目标 Q 值
        q_targets_next = self.target_q_network(next_states).detach().max(1)[0].unsqueeze(1)
        q_targets = rewards + (self.gamma * q_targets_next * (1 - dones))

        # 计算当前 Q 值
        q_expected = self.q_network(states).gather(1, actions)

        # 计算损失并更新网络
        loss = nn.MSELoss()(q_expected, q_targets)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # 软更新目标网络
        self.soft_update(self.q_network, self.target_q_network)

    def soft_update(self, local_model, target_model, tau=1e-3):
        for target_param, local_param in zip(target_model.parameters(), local_model.parameters()):
            target_param.data.copy_(tau * local_param.data + (1.0 - tau) * target_param.data)


class QNetwork(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(QNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, action_dim)

    def forward(self, state):
        x = torch.relu(self.fc1(state))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class HDRL:
    def __init__(self, state_dim, action_dim_continuous, action_dim_discrete, action_bound):
        self.ddpg = DDPG(state_dim, action_dim_continuous, action_bound)
        self.dqn = DQN(state_dim, action_dim_discrete)

    def select_action(self, state):
        continuous_action = self.ddpg.select_action(state)
        discrete_action = self.dqn.select_action(state)
        return continuous_action, discrete_action

    def update(self, batch_size):
        self.ddpg.update(batch_size)
        self.dqn.update(batch_size)

    def add_to_memory(self, state, continuous_action, discrete_action, reward, next_state, done):
        self.ddpg.memory.add(state, continuous_action, reward, next_state, done)
        self.dqn.memory.add(state, discrete_action, reward, next_state, done)


# 使用示例
# 使用示例
if __name__ == "__main__":
    state_dim = 10
    action_dim_continuous = 2
    action_dim_discrete = 5
    action_bound = 1.0

    hdrl = HDRL(state_dim, action_dim_continuous, action_dim_discrete, action_bound)
    batch_size = 64
    num_episodes = 100

    for episode in range(num_episodes):
        state = np.random.rand(state_dim)
        episode_reward = 0.0
        done = False

        while not done:
            continuous_action, discrete_action = hdrl.select_action(state)
            # 在实际应用中，这里需要与环境进行交互以获取奖励和下一个状态
            # 这里我们使用随机奖励和下一个状态进行演示
            reward = np.random.rand()
            next_state = np.random.rand(state_dim)
            done = np.random.choice([True, False], p=[0.1, 0.9])  # 随机决定是否结束episode

            hdrl.add_to_memory(state, continuous_action, discrete_action, reward, next_state, done)

            state = next_state
            episode_reward += reward

            # 当经验回放缓冲区有足够的样本时，开始训练
            if len(hdrl.ddpg.memory) >= batch_size:
                hdrl.update(batch_size)

        # 每个episode结束后打印训练信息
        print(f"Episode: {episode + 1}/{num_episodes}, Reward: {episode_reward:.4f}")
